from ultralytics import YOLO
# from com.ultralytics.models import YOLO


class TrainModel:
    """
    训练模型
    """

    def __init__(self):
        self.model = YOLO("yolo11n.pt")
        pass

    def train(self):
        self.model.train(
            data="D:\\IdentificationModel\\config\\laji.yaml",  # path to dataset YAML
            epochs=20,  # number of training epochs
            device="cpu",  # device to run on, i.e. device=0 or device=0,1,2,3 or device=cpu
        )
        # self.model.val()

    def predict(self):
        pass



if __name__ == '__main__':
    TrainModel().train()
